#include "common.h"
#include "kernel_operator.h"
#ifdef __CCE_KT_TEST__
#include <iostream>
#endif
namespace AscendC
{
template <typename _DT_X1, typename _DT_X2, typename _DT_Y> class Kernel
{
  public:
    TPipe pipe;
    CommonTiling tilingData;
    DefInTensor(X1);
    DefInTensor(X2);
    DefOutTensor(Y);

    DefBufVECIN(TEMP);
    DefBufVECIN(TEMP_1);
    DefBufVECIN(TEMP_3);

  public:
    __aicore__ inline Kernel()
    {
    }
    __aicore__ inline void Init(GM_ADDR x1, GM_ADDR x2, GM_ADDR y, CommonTiling &tiling)
    {
        tilingData = tiling;
        // 切分global
        SetGlobalBuffer(X1, x1, tilingData.startIdx, tilingData.blockLength);
        SetGlobalBuffer(X2, x2, tilingData.startIdx, tilingData.blockLength);
        SetGlobalBuffer(Y, y, tilingData.startIdx, tilingData.blockLength);
        // 初始化队列
        InitQueueSimple(X1, tilingData.bufferSize);
        InitQueueSimple(X2, tilingData.bufferSize);
        InitQueueSimple(Y, tilingData.bufferSize);
        // 初始化buf
        InitTBufBuffer(TEMP, tilingData.bufferSize);
        InitTBufBuffer(TEMP_1, tilingData.bufferSize * sizeof(half));
        InitTBufBuffer(TEMP_3, tilingData.bufferSize * sizeof(half));

        // 设置全1 temp
        TBufGet(TEMP_1, half);
        Muls(TBufTensorName(TEMP_1), TBufTensorName(TEMP_1), half(0), tilingData.bufferSize);
        Adds(TBufTensorName(TEMP_1), TBufTensorName(TEMP_1), half(1), tilingData.bufferSize);
    }
    __aicore__ inline void Process()
    {
        auto finnal_progress = tilingData.loopCount - 1;
        for (uint32_t i = 0; i < tilingData.loopCount - 1; ++i)
        {
            CopyIn(i, tilingData.bufferSize);
            Compute(i, tilingData.bufferSize);
            CopyOut(i, tilingData.bufferSize);
        }
        auto copy_count = tilingData.getFinalCopyCount<TypeOf(X1)>();
        CopyIn(finnal_progress, copy_count);
        Compute(finnal_progress, tilingData.finalCalcCount);
        CopyOut(finnal_progress, copy_count);
    }
    __aicore__ inline void CopyIn(uint32_t progress, uint32_t calcCount)
    {
        auto index = progress * tilingData.bufferSize;
        EnQueGlobal2Local(X1, index, calcCount);
        EnQueGlobal2Local(X2, index, calcCount);
    }
    __aicore__ inline void Compute(uint32_t progress, uint32_t calcCount)
    {
        DeQueSimple(X1);
        DeQueSimple(X2);
        QueAllocSimple(Y);

        TBufGet(TEMP, int8_t);
        TBufGet(TEMP_1, half);
        TBufGet(TEMP_3, half);
        Compare(TBufTensorName(TEMP), LTensorName(X1), LTensorName(X2), CMPMODE::LE, calcCount);
        Select(TBufTensorName(TEMP_3), TBufTensorName(TEMP), TBufTensorName(TEMP_1), half(0),
               SELMODE::VSEL_TENSOR_SCALAR_MODE, calcCount);
        Cast(TBufTensorName(TEMP), TBufTensorName(TEMP_3), RoundMode::CAST_TRUNC, calcCount);
        auto bool_ret = TBufTensorName(TEMP).ReinterpretCast<_DT_Y>();
        DataCopy(LTensorName(Y), bool_ret, calcCount);
        // 选取值
#ifdef __CCE_KT_TEST__
        if (progress == 0 && GetBlockIdx() == 0)
        {
            LTensorName(X1).Print(10);
            LTensorName(X2).Print(10);
            TBufTensorName(TEMP).Print(10);
            TBufTensorName(TEMP_1).Print(10);
            TBufTensorName(TEMP_3).Print(10);
            LTensorName(Y).Print(10);
        }
#endif
        EnQue(Y);
        QueFree(X1);
        QueFree(X2);
    }
    __aicore__ inline void CopyOut(uint32_t progress, uint32_t calcCount)
    {
        auto index = progress * tilingData.bufferSize;
        DeQueLocal2Global(Y, index, calcCount)
    }
};

template <typename _DT_X1, typename _DT_X2, typename _DT_Y> class KernelInt8
{
  public:
    TPipe pipe;
    CommonTiling tilingData;
    DefInTensor(X1);
    DefInTensor(X2);
    DefOutTensor(Y);

    DefBufVECIN(TEMP);
    DefBufVECIN(TEMP_1);
    DefBufVECIN(TEMP_3);

    DefBufVECIN(CASTED_X1);
    DefBufVECIN(CASTED_X2);

  public:
    __aicore__ inline KernelInt8()
    {
    }
    __aicore__ inline void Init(GM_ADDR x1, GM_ADDR x2, GM_ADDR y, CommonTiling &tiling)
    {
        tilingData = tiling;
        // 切分global
        SetGlobalBuffer(X1, x1, tilingData.startIdx, tilingData.blockLength);
        SetGlobalBuffer(X2, x2, tilingData.startIdx, tilingData.blockLength);
        SetGlobalBuffer(Y, y, tilingData.startIdx, tilingData.blockLength);
        // 初始化队列
        InitQueueSimple(X1, tilingData.bufferSize);
        InitQueueSimple(X2, tilingData.bufferSize);
        InitQueueSimple(Y, tilingData.bufferSize);
        // 初始化buf
        InitTBufBuffer(TEMP, tilingData.bufferSize);
        InitTBufBuffer(TEMP_1, tilingData.bufferSize * sizeof(half));
        InitTBufBuffer(TEMP_3, tilingData.bufferSize * sizeof(half));
        InitTBufBuffer(CASTED_X1, tilingData.bufferSize * sizeof(half));
        InitTBufBuffer(CASTED_X2, tilingData.bufferSize * sizeof(half));

        // 设置全1 temp
        TBufGet(TEMP_1, half);
        Muls(TBufTensorName(TEMP_1), TBufTensorName(TEMP_1), half(0), tilingData.bufferSize);
        Adds(TBufTensorName(TEMP_1), TBufTensorName(TEMP_1), half(1), tilingData.bufferSize);
    }
    __aicore__ inline void Process()
    {
        auto finnal_progress = tilingData.loopCount - 1;
        for (uint32_t i = 0; i < tilingData.loopCount - 1; ++i)
        {
            CopyIn(i, tilingData.bufferSize);
            Compute(i, tilingData.bufferSize);
            CopyOut(i, tilingData.bufferSize);
        }
        auto copy_count = tilingData.getFinalCopyCount<TypeOf(X1)>();
        CopyIn(finnal_progress, copy_count);
        Compute(finnal_progress, tilingData.finalCalcCount);
        CopyOut(finnal_progress, copy_count);
    }
    __aicore__ inline void CopyIn(uint32_t progress, uint32_t calcCount)
    {
        auto index = progress * tilingData.bufferSize;
        EnQueGlobal2Local(X1, index, calcCount);
        EnQueGlobal2Local(X2, index, calcCount);
    }
    __aicore__ inline void Compute(uint32_t progress, uint32_t calcCount)
    {
        DeQueSimple(X1);
        DeQueSimple(X2);
        QueAllocSimple(Y);

        TBufGet(TEMP, int8_t);
        TBufGet(TEMP_1, half);
        TBufGet(TEMP_3, half);
        TBufGet(CASTED_X1, half);
        TBufGet(CASTED_X2, half);
        Cast(TBufTensorName(CASTED_X1), LTensorName(X1), RoundMode::CAST_NONE, calcCount);
        Cast(TBufTensorName(CASTED_X2), LTensorName(X2), RoundMode::CAST_NONE, calcCount);
        Compare(TBufTensorName(TEMP), TBufTensorName(CASTED_X1), TBufTensorName(CASTED_X2), CMPMODE::LE, calcCount);
        Select(TBufTensorName(TEMP_3), TBufTensorName(TEMP), TBufTensorName(TEMP_1), half(0),
               SELMODE::VSEL_TENSOR_SCALAR_MODE, calcCount);
        Cast(TBufTensorName(TEMP), TBufTensorName(TEMP_3), RoundMode::CAST_TRUNC, calcCount);
        auto bool_ret = TBufTensorName(TEMP).ReinterpretCast<_DT_Y>();
        DataCopy(LTensorName(Y), bool_ret, calcCount);
        EnQue(Y);
        QueFree(X1);
        QueFree(X2);
    }
    __aicore__ inline void CopyOut(uint32_t progress, uint32_t calcCount)
    {
        auto index = progress * tilingData.bufferSize;
        DeQueLocal2Global(Y, index, calcCount)
    }
};

template <typename _DT_X1, typename _DT_X2, typename _DT_Y> class KernelInt32
{
  public:
    TPipe pipe;
    CommonTiling tilingData;
    DefInTensor(X1);
    DefInTensor(X2);
    DefOutTensor(Y);

    DefBufVECIN(TEMP);
    DefBufVECIN(TEMP_1);
    DefBufVECIN(TEMP_3);

    DefBufVECIN(CASTED_X1);
    DefBufVECIN(CASTED_X2);

  public:
    __aicore__ inline KernelInt32()
    {
    }
    __aicore__ inline void Init(GM_ADDR x1, GM_ADDR x2, GM_ADDR y, CommonTiling &tiling)
    {
        tilingData = tiling;
        // 切分global
        SetGlobalBuffer(X1, x1, tilingData.startIdx, tilingData.blockLength);
        SetGlobalBuffer(X2, x2, tilingData.startIdx, tilingData.blockLength);
        SetGlobalBuffer(Y, y, tilingData.startIdx, tilingData.blockLength);
        // 初始化队列
        InitQueueSimple(X1, tilingData.bufferSize);
        InitQueueSimple(X2, tilingData.bufferSize);
        InitQueueSimple(Y, tilingData.bufferSize);
        // 初始化buf
        InitTBufBuffer(TEMP, tilingData.bufferSize);
        InitTBufBuffer(TEMP_1, tilingData.bufferSize * sizeof(half));
        InitTBufBuffer(TEMP_3, tilingData.bufferSize * sizeof(half));
        InitTBufBuffer(CASTED_X1, tilingData.bufferSize * sizeof(float));
        InitTBufBuffer(CASTED_X2, tilingData.bufferSize * sizeof(float));

        // 设置全1 temp
        TBufGet(TEMP_1, half);
        Muls(TBufTensorName(TEMP_1), TBufTensorName(TEMP_1), half(0), tilingData.bufferSize);
        Adds(TBufTensorName(TEMP_1), TBufTensorName(TEMP_1), half(1), tilingData.bufferSize);
    }
    __aicore__ inline void Process()
    {
        auto finnal_progress = tilingData.loopCount - 1;
        for (uint32_t i = 0; i < tilingData.loopCount - 1; ++i)
        {
            CopyIn(i, tilingData.bufferSize);
            Compute(i, tilingData.bufferSize);
            CopyOut(i, tilingData.bufferSize);
        }
        auto copy_count = tilingData.getFinalCopyCount<TypeOf(X1)>();
        CopyIn(finnal_progress, copy_count);
        Compute(finnal_progress, tilingData.finalCalcCount);
        CopyOut(finnal_progress, copy_count);
    }
    __aicore__ inline void CopyIn(uint32_t progress, uint32_t calcCount)
    {
        auto index = progress * tilingData.bufferSize;
        EnQueGlobal2Local(X1, index, calcCount);
        EnQueGlobal2Local(X2, index, calcCount);
    }
    __aicore__ inline void Compute(uint32_t progress, uint32_t calcCount)
    {
        DeQueSimple(X1);
        DeQueSimple(X2);
        QueAllocSimple(Y);

        TBufGet(TEMP, int8_t);
        TBufGet(TEMP_1, half);
        TBufGet(TEMP_3, half);
        TBufGet(CASTED_X1, float);
        TBufGet(CASTED_X2, float);

        // 转换类型
        Cast(TBufTensorName(CASTED_X1), LTensorName(X1), RoundMode::CAST_FLOOR, calcCount);
        Cast(TBufTensorName(CASTED_X2), LTensorName(X2), RoundMode::CAST_FLOOR, calcCount);
        // 比较
        Compare(TBufTensorName(TEMP), TBufTensorName(CASTED_X1), TBufTensorName(CASTED_X2), CMPMODE::LE, calcCount);
        Select(TBufTensorName(TEMP_3), TBufTensorName(TEMP), TBufTensorName(TEMP_1), half(0),
               SELMODE::VSEL_TENSOR_SCALAR_MODE, calcCount);
        Cast(TBufTensorName(TEMP), TBufTensorName(TEMP_3), RoundMode::CAST_TRUNC, calcCount);
        auto bool_ret = TBufTensorName(TEMP).ReinterpretCast<_DT_Y>();
        DataCopy(LTensorName(Y), bool_ret, calcCount);

#ifdef __CCE_KT_TEST__
        if (progress == 0 && GetBlockIdx() == 0)
        {
            LTensorName(X1).Print(8);
            LTensorName(X2).Print(8);
            TBufTensorName(CASTED_X1).Print(8);
            TBufTensorName(CASTED_X2).Print(8);
            TBufTensorName(TEMP).Print(8);
            TBufTensorName(TEMP_1).Print(8);
            TBufTensorName(TEMP_3).Print(8);
            LTensorName(Y).Print(8);
        }
#endif
        EnQue(Y);
        QueFree(X1);
        QueFree(X2);
    }
    __aicore__ inline void CopyOut(uint32_t progress, uint32_t calcCount)
    {
        auto index = progress * tilingData.bufferSize;
        DeQueLocal2Global(Y, index, calcCount)
    }
};
} // namespace AscendC

#ifdef __CCE_KT_TEST__
extern "C" __global__ __aicore__ void less_equal(GM_ADDR x1, GM_ADDR x2, GM_ADDR y)
{
    AscendC::CommonTiling commonTiling;
    commonTiling.Init(0, 0, 1, 256, 256, 256);
    // AscendC::Kernel<half, half, bool> op;
    AscendC::KernelInt32<int32_t, int32_t, bool> op;
    op.Init(x1, x2, y, commonTiling);
    op.Process();
}
#else
extern "C" __global__ __aicore__ void less_equal(GM_ADDR x1, GM_ADDR x2, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling)
{
    GET_TILING_DATA(tiling_data, tiling);
    CRATE_COMMON_TILING_SIMPLE;
    switch (tiling_data.dType)
    {
    // float
    case 0: {
        AscendC::Kernel<float, float, DTYPE_Y> op;
        op.Init(x1, x2, y, commonTiling);
        op.Process();
    }
    break;
    // half
    case 1: {
        AscendC::Kernel<half, half, DTYPE_Y> op;
        op.Init(x1, x2, y, commonTiling);
        op.Process();
    }
    break;
    // int8_t
    case 2: {
        AscendC::KernelInt8<int8_t, int8_t, DTYPE_Y> op;
        op.Init(x1, x2, y, commonTiling);
        op.Process();
    }
    break;
    // int32_t
    case 3: {
        AscendC::KernelInt32<int32_t, int32_t, DTYPE_Y> op;
        op.Init(x1, x2, y, commonTiling);
        op.Process();
    }
    break;
    }
}
#endif
